{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "772ebb07-8c83-4e20-97d4-8d68ef2d642f",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/workflow/human_in_the_loop_story_crafting.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4a4fc96-58a8-4b0b-a4fb-077ef16b59e2",
   "metadata": {},
   "source": [
    "# Choose Your Own Adventure Workflow (Human In The Loop)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e631036a-1a21-4f5b-a497-b95f8a9438f5",
   "metadata": {},
   "source": [
    "For some Workflow applications, it may desirable and/or required to have humans involved in its execution. For example, a step of a Workflow may need human expertise or input in order to run. In another scenario, it may be required to have a human validate the initial output of a Workflow.\n",
    "\n",
    "In this notebook, we show how one can implement a human-in-the-loop pattern with Workflows. Here we'll build a Workflow that creates stories in the style of Choose Your Own Adventure, where the LLM produces a segment of the story along with potential actions, and a human is required to choose from one of those actions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b69b368-c7f7-4aee-9e15-4fcd325e238a",
   "metadata": {},
   "source": [
    "## Generating Segments Of The Story With An LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f093ac20-386d-4eac-86c5-0d4134a77099",
   "metadata": {},
   "source": [
    "Here, we'll make use of the ability to produce structured outputs from an LLM. We will task the LLM to create a segment of the story that is in continuation of previously generated segments and action choices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e389bf9f-3048-4372-b967-9aa909990eef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, List\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core.bridge.pydantic import BaseModel, Field\n",
    "from llama_index.core.prompts import PromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d42111a0-cb06-47f1-af60-e4e4628e8c16",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Segment(BaseModel):\n",
    "    \"\"\"Data model for generating segments of a story.\"\"\"\n",
    "\n",
    "    plot: str = Field(\n",
    "        description=\"The plot of the adventure for the current segment. The plot should be no longer than 3 sentences.\"\n",
    "    )\n",
    "    actions: List[str] = Field(\n",
    "        default=[],\n",
    "        description=\"The list of actions the protaganist can take that will shape the plot and actions of the next segment.\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6826965f-58eb-4027-a431-069734ed671f",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEGMENT_GENERATION_TEMPLATE = \"\"\"\n",
    "You are working with a human to create a story in the style of choose your own adventure.\n",
    "\n",
    "The human is playing the role of the protaganist in the story which you are tasked to\n",
    "help write. To create the story, we do it in steps, where each step produces a BLOCK.\n",
    "Each BLOCK consists of a PLOT, a set of ACTIONS that the protaganist can take, and the\n",
    "chosen ACTION. \n",
    "\n",
    "Below we attach the history of the adventure so far.\n",
    "\n",
    "PREVIOUS BLOCKS:\n",
    "---\n",
    "{running_story}\n",
    "\n",
    "Continue the story by generating the next block's PLOT and set of ACTIONs. If there are\n",
    "no previous BLOCKs, start an interesting brand new story. Give the protaganist a name and an\n",
    "interesting challenge to solve.\n",
    "\n",
    "\n",
    "Use the provided data model to structure your output.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "758b5b54-c342-4fd0-b1b4-a58a54707a86",
   "metadata": {},
   "outputs": [],
   "source": [
    "FINAL_SEGMENT_GENERATION_TEMPLATE = \"\"\"\n",
    "You are working with a human to create a story in the style of choose your own adventure.\n",
    "\n",
    "The human is playing the role of the protaganist in the story which you are tasked to\n",
    "help write. To create the story, we do it in steps, where each step produces a BLOCK.\n",
    "Each BLOCK consists of a PLOT, a set of ACTIONS that the protaganist can take, and the\n",
    "chosen ACTION. Below we attach the history of the adventure so far.\n",
    "\n",
    "PREVIOUS BLOCKS:\n",
    "---\n",
    "{running_story}\n",
    "\n",
    "The story is now coming to an end. With the previous blocks, wrap up the story with a\n",
    "closing PLOT. Since it is a closing plot, DO NOT GENERATE a new set of actions.\n",
    "\n",
    "Use the provided data model to structure your output.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f92da9d-e2e2-4998-a019-bb95e36c9f3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's see an example segment\n",
    "llm = OpenAI(\"gpt-4o\")\n",
    "segment = llm.structured_predict(\n",
    "    Segment,\n",
    "    PromptTemplate(SEGMENT_GENERATION_TEMPLATE),\n",
    "    running_story=\"\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c51333c-d95f-4127-97ae-e69178988ff2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Segment(plot=\"In the bustling city of Eldoria, a young adventurer named Aric discovered a mysterious map hidden inside an old bookshop. The map hinted at a hidden treasure buried deep within the enchanted Whispering Woods. Intrigued and eager for adventure, Aric decided to follow the map's clues.\", actions=['Follow the map to the Whispering Woods', 'Seek advice from the old bookshop owner', 'Gather supplies for the journey', 'Ignore the map and continue with daily life'])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "segment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1408ab3-3764-415a-b70c-d2fca69a2dac",
   "metadata": {},
   "source": [
    "### Stitching together previous segments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfbdad71-ff46-4044-b2a6-68721f32df81",
   "metadata": {},
   "source": [
    "We need to stich together story segments and pass this in to the prompt as the value for `running_story`. We define a `Block` data class that holds the `Segment` as well as the `choice` of action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19c4ed39-a2af-4c89-9dba-ddd49aaa0cd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "from typing import Optional\n",
    "\n",
    "BLOCK_TEMPLATE = \"\"\"\n",
    "BLOCK\n",
    "===\n",
    "PLOT: {plot}\n",
    "ACTIONS: {actions}\n",
    "CHOICE: {choice}\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "class Block(BaseModel):\n",
    "    id_: str = Field(default_factory=lambda: str(uuid.uuid4()))\n",
    "    segment: Segment\n",
    "    choice: Optional[str] = None\n",
    "    block_template: str = BLOCK_TEMPLATE\n",
    "\n",
    "    def __str__(self):\n",
    "        return self.block_template.format(\n",
    "            plot=self.segment.plot,\n",
    "            actions=\", \".join(self.segment.actions),\n",
    "            choice=self.choice or \"\",\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dda2aed4-9df1-4d00-813c-8298b39c5e85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "BLOCK\n",
      "===\n",
      "PLOT: In the bustling city of Eldoria, a young adventurer named Aric discovered a mysterious map hidden inside an old bookshop. The map hinted at a hidden treasure buried deep within the enchanted Whispering Woods. Intrigued and eager for adventure, Aric decided to follow the map's clues.\n",
      "ACTIONS: Follow the map to the Whispering Woods, Seek advice from the old bookshop owner, Gather supplies for the journey, Ignore the map and continue with daily life\n",
      "CHOICE: \n",
      "\n"
     ]
    }
   ],
   "source": [
    "block = Block(segment=segment)\n",
    "print(block)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c103bff-342c-4b88-9f34-01489ec2dc76",
   "metadata": {},
   "source": [
    "## Create The Choose Your Own Adventure Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76ae34dd-9ce2-42d4-82ca-e6525fc163c3",
   "metadata": {},
   "source": [
    "This Workflow will consist of two steps that will cycle until a max number of steps (i.e., segments) has been produced. The first step will have the LLM create a new `Segment`, which will be used to create a new story `Block`. The second step will prompt the human to choose their adventure from the list of actions specified in the newly created `Segment`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf4d468a-f4de-478c-9f7b-8931d4dad4ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.workflow import (\n",
    "    Context,\n",
    "    Event,\n",
    "    StartEvent,\n",
    "    StopEvent,\n",
    "    Workflow,\n",
    "    step,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abed042e-afb7-4e29-9d23-1374d3583c30",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NewBlockEvent(Event):\n",
    "    block: Block\n",
    "\n",
    "\n",
    "class HumanChoiceEvent(Event):\n",
    "    block_id: str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83ed549a-36ec-446c-963f-d0ced185622f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChooseYourOwnAdventureWorkflow(Workflow):\n",
    "    def __init__(self, max_steps: int = 3, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.llm = OpenAI(\"gpt-4o\")\n",
    "        self.max_steps = max_steps\n",
    "\n",
    "    @step\n",
    "    async def create_segment(\n",
    "        self, ctx: Context, ev: StartEvent | HumanChoiceEvent\n",
    "    ) -> NewBlockEvent | StopEvent:\n",
    "        blocks = await ctx.store.get(\"blocks\", [])\n",
    "        running_story = \"\\n\".join(str(b) for b in blocks)\n",
    "\n",
    "        if len(blocks) < self.max_steps:\n",
    "            new_segment = self.llm.structured_predict(\n",
    "                Segment,\n",
    "                PromptTemplate(SEGMENT_GENERATION_TEMPLATE),\n",
    "                running_story=running_story,\n",
    "            )\n",
    "            new_block = Block(segment=new_segment)\n",
    "            blocks.append(new_block)\n",
    "            await ctx.store.set(\"blocks\", blocks)\n",
    "            return NewBlockEvent(block=new_block)\n",
    "        else:\n",
    "            final_segment = self.llm.structured_predict(\n",
    "                Segment,\n",
    "                PromptTemplate(FINAL_SEGMENT_GENERATION_TEMPLATE),\n",
    "                running_story=running_story,\n",
    "            )\n",
    "            final_block = Block(segment=final_segment)\n",
    "            blocks.append(final_block)\n",
    "            return StopEvent(result=blocks)\n",
    "\n",
    "    @step\n",
    "    async def prompt_human(\n",
    "        self, ctx: Context, ev: NewBlockEvent\n",
    "    ) -> HumanChoiceEvent:\n",
    "        block = ev.block\n",
    "\n",
    "        # get human input\n",
    "        human_prompt = f\"\\n===\\n{ev.block.segment.plot}\\n\\n\"\n",
    "        human_prompt += \"Choose your adventure:\\n\\n\"\n",
    "        human_prompt += \"\\n\".join(ev.block.segment.actions)\n",
    "        human_prompt += \"\\n\\n\"\n",
    "        human_input = input(human_prompt)\n",
    "\n",
    "        blocks = await ctx.store.get(\"blocks\")\n",
    "        block.choice = human_input\n",
    "        blocks[-1] = block\n",
    "        await ctx.store.set(\"block\", blocks)\n",
    "\n",
    "        return HumanChoiceEvent(block_id=ev.block.id_)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "533cc3c8-d29a-4448-9237-a45855628d1e",
   "metadata": {},
   "source": [
    "### Running The Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "431dc17c-16cc-4adb-a2d8-c9ee117e6480",
   "metadata": {},
   "source": [
    "Since workflows are async first, this all runs fine in a notebook. If you were running in your own code, you would want to use `asyncio.run()` to start an async event loop if one isn't already running.\n",
    "\n",
    "```python\n",
    "async def main():\n",
    "    <async code>\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    import asyncio\n",
    "    asyncio.run(main())\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "260e3f7b-a490-4708-b589-4e829a1e2580",
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07fcdf83-b2ee-43ab-b2df-61931aca2888",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = ChooseYourOwnAdventureWorkflow(timeout=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba835938-d5cf-489f-b27a-0a0f76119cee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdin",
     "output_type": "stream",
     "text": [
      "\n",
      "===\n",
      "In the bustling city of Eldoria, a young adventurer named Aric discovered a mysterious map hidden in an old bookshop. The map hinted at a hidden treasure buried deep within the Whispering Woods, a place known for its eerie silence and ancient secrets. Determined to uncover the treasure, Aric set off on his journey, leaving the city behind.\n",
      "\n",
      "Choose your adventure:\n",
      "\n",
      "Follow the map into the Whispering Woods\n",
      "Seek help from a local guide\n",
      "Study the map for hidden clues\n",
      "\n",
      " Seek help from a local guide\n",
      "\n",
      "===\n",
      "Aric found a seasoned guide named Elara, who knew the Whispering Woods like the back of her hand. Elara agreed to help Aric, intrigued by the promise of hidden treasure. Together, they ventured into the forest, the map leading them to a fork in the path where the trees seemed to whisper secrets.\n",
      "\n",
      "Choose your adventure:\n",
      "\n",
      "Take the left path\n",
      "Take the right path\n",
      "Ask Elara for advice\n",
      "\n",
      " Ask Elara for advice\n",
      "\n",
      "===\n",
      "Elara examined the map closely and listened to the whispers of the trees. She suggested taking the left path, as it seemed to align with the ancient markings on the map. Trusting her expertise, Aric and Elara proceeded down the left path, where the forest grew denser and the air filled with an eerie stillness.\n",
      "\n",
      "Choose your adventure:\n",
      "\n",
      "Continue down the left path\n",
      "Turn back and take the right path\n",
      "Look for hidden clues along the path\n",
      "\n",
      " Look for hidden clues along the path\n"
     ]
    }
   ],
   "source": [
    "result = await w.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6c29727-1614-4a8c-b56b-69b5f9b553e1",
   "metadata": {},
   "source": [
    "### Print The Final Story"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e15aa0c-78da-4ff6-9519-d0b134bdd546",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In the bustling city of Eldoria, a young adventurer named Aric discovered a mysterious map hidden in an old bookshop. The map hinted at a hidden treasure buried deep within the Whispering Woods, a place known for its eerie silence and ancient secrets. Determined to uncover the treasure, Aric set off on his journey, leaving the city behind.\n",
      "\n",
      "Aric found a seasoned guide named Elara, who knew the Whispering Woods like the back of her hand. Elara agreed to help Aric, intrigued by the promise of hidden treasure. Together, they ventured into the forest, the map leading them to a fork in the path where the trees seemed to whisper secrets.\n",
      "\n",
      "Elara examined the map closely and listened to the whispers of the trees. She suggested taking the left path, as it seemed to align with the ancient markings on the map. Trusting her expertise, Aric and Elara proceeded down the left path, where the forest grew denser and the air filled with an eerie stillness.\n",
      "\n",
      "As Aric and Elara searched for hidden clues along the dense path, they stumbled upon an ancient stone altar covered in moss and vines. Upon closer inspection, they discovered a hidden compartment within the altar containing the long-lost treasure—a chest filled with gold, jewels, and ancient artifacts. With their mission complete, Aric and Elara returned to Eldoria, their bond strengthened by the adventure and their hearts filled with the thrill of discovery.\n"
     ]
    }
   ],
   "source": [
    "final_story = \"\\n\\n\".join(b.segment.plot for b in result)\n",
    "print(final_story)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ce77433-8e18-4793-a7bb-604f5ba0c86c",
   "metadata": {},
   "source": [
    "### Other Ways To Implement Human In The Loop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20465b96-9382-4a4c-a2ec-767a092ba068",
   "metadata": {},
   "source": [
    "One could also implement the human in the loop by creating a separate Workflow just for gathering human input and making use of nested Workflows. This design could be used in situations where you would want the human input gathering to be a separate service from the rest of the Workflow, which is what would happen if you deployed the nested workflows with llama-deploy."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index-core",
   "language": "python",
   "name": "llama-index-core"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
